package eval

import (
	"fmt"
	"showen/ast"
	"showen/obj"
	"showen/tok"
	"sync/atomic"
)

// A small set of language keywords and core procedures that should not be renamed
// by hygiene (they are part of the language's surface syntax or globals).
var coreLiterals = map[string]bool{
	"if":     true,
	"begin":  true,
	"set!":   true,
	"let":    true,
	"lambda": true,
	"quote":  true,
	"not":    true,
}

// applyDefineSyntax implements define-syntax with syntax-rules.
// It supports literal identifiers, pattern variables, and a limited
// ellipsis implementation (ellipsis only allowed at the end of a pattern
// subform, which covers common macros like `when`/`unless`).
func applyDefineSyntax(args []ast.Exp, env *obj.Environment) (obj.Object, error) {
	if len(args) != 2 {
		return nil, fmt.Errorf("define-syntax requires exactly 2 arguments")
	}

	nameIdent, ok := args[0].(*ast.IdentifierExp)
	if !ok {
		return nil, fmt.Errorf("first argument to define-syntax must be an identifier")
	}

	srExp, ok := args[1].(*ast.ListExp)
	if !ok {
		return nil, fmt.Errorf("second argument to define-syntax must be a syntax-rules form")
	}

	// srExp should be (syntax-rules (literals...) clause...)
	if len(srExp.Exps) < 2 {
		return nil, fmt.Errorf("invalid syntax-rules form")
	}

	headIdent, ok := srExp.Exps[0].(*ast.IdentifierExp)
	if !ok || headIdent.Token.Lexeme != "syntax-rules" {
		return nil, fmt.Errorf("expected syntax-rules")
	}

	// literals list
	literalsList, ok := srExp.Exps[1].(*ast.ListExp)
	if !ok {
		return nil, fmt.Errorf("syntax-rules requires a list of literal identifiers")
	}

	literals := []string{}
	for _, lit := range literalsList.Exps {
		if id, ok := lit.(*ast.IdentifierExp); ok {
			literals = append(literals, id.Token.Lexeme)
		} else {
			return nil, fmt.Errorf("literal list must contain identifiers")
		}
	}

	// The macro name itself is implicitly a literal in patterns
	literals = append(literals, nameIdent.Token.Lexeme)

	// Parse clauses
	var clauses []obj.MacroClause
	for _, c := range srExp.Exps[2:] {
		clauseList, ok := c.(*ast.ListExp)
		if !ok || len(clauseList.Exps) < 2 {
			return nil, fmt.Errorf("each clause must be a list with a pattern and a template")
		}
		pattern := clauseList.Exps[0]
		// Template may be multiple expressions; if so, wrap in begin
		var template ast.Exp
		if len(clauseList.Exps[1:]) == 1 {
			template = clauseList.Exps[1]
		} else {
			template = &ast.ListExp{Exps: append([]ast.Exp{&ast.IdentifierExp{Token: tok.Token{Type: tok.IDENTIFIER, Lexeme: "begin", Literal: "begin"}}}, clauseList.Exps[1:]...)}
		}
		clauses = append(clauses, obj.MacroClause{Pattern: pattern, Template: template})
	}

	macro := &obj.MacroObject{Name: nameIdent.Token.Lexeme, Literals: literals, Clauses: clauses, DefEnv: env}
	env.Set(nameIdent.Token.Lexeme, macro)
	return obj.NULL_OBJECT, nil
}

// expandMacro attempts to expand a macro call (`call` is the full application
// list expression) using the macro's clauses. It returns an AST expression
// ready for evaluation.
func expandMacro(m *obj.MacroObject, call *ast.ListExp) (ast.Exp, error) {
	// For each clause, try to match
	for _, cl := range m.Clauses {
		env := newMatchEnv()
		if matchPattern(cl.Pattern, call, m.Literals, env) {
			// Perform template substitution with hygiene (per-expansion renaming)
			rename := map[string]string{}
			res, err := substituteWithRename(cl.Template, env, rename, m.Literals, m.DefEnv)
			if err != nil {
				// Wrap error with macro name and clause/pattern location for easier debugging
				patStr := "<unknown>"
				if cl.Pattern != nil {
					patStr = cl.Pattern.String()
				}
				return nil, fmt.Errorf("macro %s, clause pattern %s: %w", m.Name, patStr, err)
			}
			return res, nil
		}
	}
	return nil, fmt.Errorf("no matching syntax-rules clause for macro %s", m.Name)
}

// matchEnv stores pattern variable bindings. For simple variables we store
// a single ast.Exp; for repeated variables we store a slice.
// matchEnv stores pattern variable bindings.
// single: simple single bindings for identifiers
// repeat: for a given variable name, a slice of *matchEnv where each element
// represents the bindings that occurred for one repetition of the repeated
// subpattern. The per-repetition *matchEnv may itself contain nested repeats
// in its repeat map, enabling nested grouped captures.
type matchEnv struct {
	single map[string]ast.Exp
	repeat map[string][]*matchEnv
}

func newMatchEnv() *matchEnv {
	return &matchEnv{single: map[string]ast.Exp{}, repeat: map[string][]*matchEnv{}}
}

func isEllipsisId(e ast.Exp) bool {
	if id, ok := e.(*ast.IdentifierExp); ok {
		return id.Token.Lexeme == "..."
	}
	return false
}

// matchPattern supports matching of list patterns. It handles pattern variables,
// literals, and a limited form of ellipsis: if a pattern element is followed by
// an identifier with lexeme "..." and it's at the end of the pattern sublist,
// then the corresponding datum tail is bound to that variable (zero or more matches).
func matchPattern(pat ast.Exp, datum ast.Exp, literals []string, env *matchEnv) bool {
	// For lists, delegate to sequence matcher which supports ellipsis at arbitrary positions
	if _, ok := pat.(*ast.ListExp); ok {
		pList, _ := pat.(*ast.ListExp)
		if dList, ok2 := datum.(*ast.ListExp); ok2 {
			return matchSeq(pList.Exps, dList.Exps, literals, env)
		}
		return false
	}

	// Non-list pattern: identifier or literal
	if pid, ok := pat.(*ast.IdentifierExp); ok {
		// underscore wildcard
		if pid.Token.Lexeme == "_" {
			return true
		}
		// literal?
		for _, lit := range literals {
			if pid.Token.Lexeme == lit {
				if did, ok := datum.(*ast.IdentifierExp); ok {
					return did.Token.Lexeme == pid.Token.Lexeme
				}
				return false
			}
		}
		// bind as single
		env.single[pid.Token.Lexeme] = datum
		return true
	}
	return false
}

// matchSeq attempts to match a sequence of pattern elements pExps against data elements dExps.
// It supports ellipsis after any pattern element with backtracking. Repeated matches
// collect per-iteration bindings into env.repeat (as slices of ast.Exp for each variable).
func matchSeq(pExps []ast.Exp, dExps []ast.Exp, literals []string, env *matchEnv) bool {
	// helper to compute minimal data length needed to match remaining patterns
	minNeeded := func(pats []ast.Exp) int {
		n := 0
		for i := 0; i < len(pats); i++ {
			if i+1 < len(pats) && isEllipsisId(pats[i+1]) {
				// repetition can match zero times
				i++
				continue
			}
			n++
		}
		return n
	}

	var backtrack func(pi, di int) bool
	// snapshot helpers
	snapshot := func() (map[string]ast.Exp, map[string][]*matchEnv) {
		s1 := make(map[string]ast.Exp)
		for k, v := range env.single {
			s1[k] = v
		}
		s2 := make(map[string][]*matchEnv)
		for k, v := range env.repeat {
			// shallow copy slice
			copied := make([]*matchEnv, len(v))
			copy(copied, v)
			s2[k] = copied
		}
		return s1, s2
	}
	restore := func(s1 map[string]ast.Exp, s2 map[string][]*matchEnv) {
		env.single = s1
		env.repeat = s2
	}

	backtrack = func(pi, di int) bool {
		if pi == len(pExps) {
			return di == len(dExps)
		}
		// If next is ellipsis, try repetitions
		if pi+1 < len(pExps) && isEllipsisId(pExps[pi+1]) {
			// Max repetitions constrained by remaining data and minimal requirement for the rest
			minRem := minNeeded(pExps[pi+2:])
			maxK := len(dExps) - di - minRem
			if maxK < 0 {
				return false
			}
			// collect variable names that appear inside the repeated subpattern
			namesInPat := map[string]bool{}
			collectIdentifiers(pExps[pi], namesInPat)
			// remove ellipsis and literal names
			for lit := range namesInPat {
				if lit == "..." {
					delete(namesInPat, lit)
				}
			}

			for k := 0; k <= maxK; k++ {
				// try match k repetitions
				s1, s2 := snapshot()
				ok := true
				// For each repetition, match pExps[pi] against dExps[di + j]
				for j := 0; j < k; j++ {
					// collect bindings from single match without mutating env (use temp env)
					temp := newMatchEnv()
					if !matchPatternCollect(pExps[pi], dExps[di+j], literals, temp) {
						ok = false
						break
					}
					// For this repetition, append the repetition-env (temp) for each
					// identifier name that occurs in the repeated pattern. This ensures
					// placeholders are produced even when a repetition yields zero
					// bindings for that name (so nested-empty lists are preserved).
					for name := range namesInPat {
						env.repeat[name] = append(env.repeat[name], temp)
					}
				}
				if ok {
					if backtrack(pi+2, di+k) {
						return true
					}
				}
				restore(s1, s2)
			}
			return false
		}

		// Normal single-element match
		if di >= len(dExps) {
			return false
		}
		// collect bindings into temp, then merge into env.single
		temp := newMatchEnv()
		if !matchPatternCollect(pExps[pi], dExps[di], literals, temp) {
			return false
		}
		// merge temp.single into env.single
		for name, val := range temp.single {
			env.single[name] = val
		}
		// merge any inner-repeat entries into the parent's env.repeat so that
		// non-repeated contexts expose inner repeats globally for substitution
		for name, reps := range temp.repeat {
			env.repeat[name] = append(env.repeat[name], reps...)
		}
		return backtrack(pi+1, di+1)
	}

	return backtrack(0, 0)
}

// matchPatternCollect matches a single pattern element against a datum, collecting
// bindings into the provided matchEnv without mutating the caller's env. It returns
// true if matched.
func matchPatternCollect(pat ast.Exp, datum ast.Exp, literals []string, env *matchEnv) bool {
	// Non-list
	if pid, ok := pat.(*ast.IdentifierExp); ok {
		if pid.Token.Lexeme == "_" {
			return true
		}
		for _, lit := range literals {
			if pid.Token.Lexeme == lit {
				if did, ok := datum.(*ast.IdentifierExp); ok {
					return did.Token.Lexeme == pid.Token.Lexeme
				}
				return false
			}
		}
		env.single[pid.Token.Lexeme] = datum
		return true
	}
	if pList, ok := pat.(*ast.ListExp); ok {
		dList, ok2 := datum.(*ast.ListExp)
		if !ok2 {
			return false
		}
		return matchSeq(pList.Exps, dList.Exps, literals, env)
	}
	return false
}

// substitute applies the environment bindings to the template, producing a new AST.
// It handles simple variable substitution and the limited ellipsis expansion: when
// the template contains (var ... ) at list positions, the var's repeated bindings
// are spliced into the result.
// gensym counter for unique names
var gensymCounter uint64

func gensym(base string) string {
	n := atomic.AddUint64(&gensymCounter, 1)
	return fmt.Sprintf("__%s_%d", base, n)
}

// substituteWithRename applies the environment bindings to the template with
// per-expansion renaming (hygiene). rename is a map from original identifier
// names in the template to generated fresh names for this expansion.
// objectToAst converts an obj.Object (from definition environment) into an AST
// expression that can be spliced into the template.
func objectToAst(o obj.Object) (ast.Exp, error) {
	switch v := o.(type) {
	case *obj.BooleanObject:
		lex := "#f"
		if v.Value {
			lex = "#t"
		}
		return &ast.BooleanExp{Token: tok.Token{Type: tok.BOOLEAN, Lexeme: lex, Literal: v.Value}}, nil
	case *obj.SymbolObject:
		return &ast.IdentifierExp{Token: tok.Token{Type: tok.IDENTIFIER, Lexeme: v.Value, Literal: v.Value}}, nil
	case *obj.StringObject:
		return &ast.StringExp{Token: tok.Token{Type: tok.STRING, Lexeme: fmt.Sprintf("\"%s\"", v.Value), Literal: v.Value}}, nil
	case *obj.NumberObject:
		return &ast.NumberExp{Token: tok.Token{Type: tok.INTEGER, Lexeme: v.Value.String(), Literal: v.Value}}, nil
	case *obj.PairObject:
		// Build list
		var exps []ast.Exp
		current := v
		for {
			// car
			if carObj, ok := current.Car.(obj.Object); ok {
				ae, err := objectToAst(carObj)
				if err != nil {
					return nil, err
				}
				exps = append(exps, ae)
			} else {
				return nil, fmt.Errorf("unsupported pair car type: %T", current.Car)
			}
			// cdr
			switch cdr := current.Cdr.(type) {
			case *obj.PairObject:
				current = cdr
				continue
			case *obj.EmptyList:
				return &ast.ListExp{Exps: exps}, nil
			case obj.Object:
				ae, err := objectToAst(cdr)
				if err != nil {
					return nil, err
				}
				exps = append(exps, ae)
				return &ast.ListExp{Exps: exps}, nil
			default:
				return nil, fmt.Errorf("unsupported pair cdr type: %T", current.Cdr)
			}
		}
	case *obj.VectorObject:
		var elems []ast.Exp
		for _, e := range v.Elements {
			ae, err := objectToAst(e)
			if err != nil {
				return nil, err
			}
			elems = append(elems, ae)
		}
		return &ast.VectorExp{Exps: elems}, nil
	default:
		return nil, fmt.Errorf("unsupported object to ast conversion: %T", o)
	}
}

// substituteWithRename is a convenience wrapper that starts substitution with an
// empty repetition context.
func substituteWithRename(t ast.Exp, env *matchEnv, rename map[string]string, literals []string, defEnv *obj.Environment) (ast.Exp, error) {
	return substituteWithRenameCtx(t, env, rename, literals, defEnv, nil)
}

// findRepsForExp walks an expression to locate an identifier which has a repetition
// slice available either in the local context (ctx) or the global env. It returns
// the slice of *matchEnv representing the repetitions for that identifier, or nil.
func findRepsForExp(e ast.Exp, env *matchEnv, ctx *matchEnv) []*matchEnv {
	switch ex := e.(type) {
	case *ast.IdentifierExp:
		if ctx != nil {
			if reps, ok := ctx.repeat[ex.Token.Lexeme]; ok {
				return reps
			}
		}
		if reps, ok := env.repeat[ex.Token.Lexeme]; ok {
			return reps
		}
		return nil
	case *ast.ListExp:
		for _, child := range ex.Exps {
			if r := findRepsForExp(child, env, ctx); r != nil {
				return r
			}
		}
		return nil
	default:
		return nil
	}
}

// collectIdentifiers returns the set of identifier names found in an expression
// (walks lists recursively).
func collectIdentifiers(e ast.Exp, out map[string]bool) {
	switch ex := e.(type) {
	case *ast.IdentifierExp:
		out[ex.Token.Lexeme] = true
	case *ast.ListExp:
		for _, c := range ex.Exps {
			collectIdentifiers(c, out)
		}
	}
}

// isPatternVar returns true if the given name appears as a bound pattern variable
// in the provided match environments (either top-level env or ctx).
func isPatternVar(name string, env *matchEnv, ctx *matchEnv) bool {
	if env != nil {
		if _, ok := env.single[name]; ok {
			return true
		}
		if _, ok := env.repeat[name]; ok {
			return true
		}
	}
	if ctx != nil {
		if _, ok := ctx.single[name]; ok {
			return true
		}
		if _, ok := ctx.repeat[name]; ok {
			return true
		}
	}
	return false
}

func quoteIfList(e ast.Exp) ast.Exp {
	switch v := e.(type) {
	case *ast.ListExp, *ast.VectorExp:
		return &ast.QuoteExp{Token: tok.Token{Type: tok.QUOTE, Lexeme: "'"}, Value: v}
	default:
		return e
	}
}

// substituteWithRenameCtx applies bindings with hygiene while using ctx as the
// current per-repetition environment (or nil at top-level). ctx allows nested
// repetitions to resolve to per-repetition bindings.
func substituteWithRenameCtx(t ast.Exp, env *matchEnv, rename map[string]string, literals []string, defEnv *obj.Environment, ctx *matchEnv) (ast.Exp, error) {
	switch te := t.(type) {
	case *ast.IdentifierExp:
		// check current repetition context first
		if ctx != nil {
			if v, ok := ctx.single[te.Token.Lexeme]; ok {
				return quoteIfList(v), nil
			}
		}
		// global single
		if v, ok := env.single[te.Token.Lexeme]; ok {
			return quoteIfList(v), nil
		}
		// if bound as repeated but used in non-ellipsis position, return first if any
		if reps, ok := env.repeat[te.Token.Lexeme]; ok {
			if len(reps) > 0 {
				rep0 := reps[0]
				if v, ok := rep0.single[te.Token.Lexeme]; ok {
					return quoteIfList(v), nil
				}
			}
		}
		// If it's a literal (from syntax-rules) or a core language symbol, keep as-is
		for _, lit := range literals {
			if te.Token.Lexeme == lit {
				return te, nil
			}
		}
		if coreLiterals[te.Token.Lexeme] {
			return te, nil
		}
		// If this identifier names a binding in the macro's defining environment,
		// splice the definition-time binding into the template (as a literal AST)
		if defEnv != nil {
			if val, ok := defEnv.Get(te.Token.Lexeme); ok {
				if ae, err := objectToAst(val); err == nil {
					return ae, nil
				}
				// If conversion fails, treat this identifier as a definition-time
				// reference and leave it as the original identifier (don't hygiene-rename).
				return te, nil
			}
		}
		// Otherwise, generate a fresh name for hygiene, unless this identifier is
		// actually a pattern variable present in the match environment (either
		// top-level env or the current ctx); in that case leave it unchanged.
		if isPatternVar(te.Token.Lexeme, env, ctx) {
			// If it's a pattern variable but there is no binding in the current
			// repetition context or as a single binding, treat it as empty (no
			// generated expression) rather than leaving an unbound identifier.
			if ctx != nil {
				if _, ok := ctx.single[te.Token.Lexeme]; !ok {
					// No binding in this repetition -> no expression
					return nil, nil
				}
			}
			if _, ok := env.single[te.Token.Lexeme]; !ok {
				// No top-level single binding -> no expression
				return nil, nil
			}
			// Otherwise fallthrough to renaming/returning the bound single below
		}
		if r, ok := rename[te.Token.Lexeme]; ok {
			return &ast.IdentifierExp{Token: tok.Token{Type: tok.IDENTIFIER, Lexeme: r, Literal: r}}, nil
		}
		newname := gensym(te.Token.Lexeme)
		rename[te.Token.Lexeme] = newname
		return &ast.IdentifierExp{Token: tok.Token{Type: tok.IDENTIFIER, Lexeme: newname, Literal: newname}}, nil
	case *ast.ListExp:
		var out []ast.Exp
		i := 0
		for i < len(te.Exps) {
			// If next is ellipsis, expand repetitions for this subexpression
			if i+1 < len(te.Exps) && isEllipsisId(te.Exps[i+1]) {
				reps := findRepsForExp(te.Exps[i], env, ctx)
				// If we have an explicit ctx (a per-repetition environment) and it
				// contains an entry for the repeated identifiers, honor it even if
				// it's empty (zero repetitions). Only when there is no ctx do we
				// attempt the identifier-scanning fallback to find global repetitions.
				if len(reps) == 0 && ctx == nil {
					ids := map[string]bool{}
					collectIdentifiers(te.Exps[i], ids)
					for id := range ids {
						if r, ok := env.repeat[id]; ok {
							reps = r
							break
						}
					}
				}
				if len(reps) > 0 {
					// stricter validation: ensure that repetition groups referenced in
					// this subexpression are consistent in length. Collect identifiers
					// inside the repeated subexpression and check their repetition lengths.
					ids := map[string]bool{}
					collectIdentifiers(te.Exps[i], ids)
					var expectedLen int = -1
					var inconsistent []string
					for id := range ids {
						if r, ok := env.repeat[id]; ok {
							if expectedLen == -1 {
								expectedLen = len(r)
							} else if len(r) != expectedLen {
								inconsistent = append(inconsistent, id)
							}
						}
					}
					if expectedLen != -1 && len(inconsistent) > 0 {
						return nil, fmt.Errorf("inconsistent repetition lengths in template for ids: %v", inconsistent)
					}
					for _, repEnv := range reps {
						r, err := substituteWithRenameCtx(te.Exps[i], env, rename, literals, defEnv, repEnv)
						if err != nil {
							return nil, err
						}
						if r != nil {
							out = append(out, r)
						}
					}
					i += 2
					continue
				}
				// No repetition bindings found; skip expansion
				i += 2
				continue
			}
			r, err := substituteWithRenameCtx(te.Exps[i], env, rename, literals, defEnv, ctx)
			if err != nil {
				return nil, err
			}
			out = append(out, r)
			i++
		}
		return &ast.ListExp{Exps: out}, nil
	default:
		return te, nil
	}
}
